{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJcYs_ERTnnI"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMUDt0CiUJk9"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77z2OchJTk0l"
      },
      "source": [
        "# Migration Examples: Canned Estimators\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/canned_estimators\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/canned_estimators.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/canned_estimators.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/canned_estimators.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUTrR4I6m1C"
      },
      "source": [
        "Canned (or Premade) Estimators have traditionally been used in TensorFlow 1 as quick and easy ways to train models for a variety of typical use cases. TensorFlow 2 provides straightforward approximate substitutes for a number of them by way of Keras models. For those canned estimators that do not have built-in TF2 substitutes, you can still build your own replacement fairly easily.\n",
        "\n",
        "This guide walks through a few examples of direct equivalents and custom substitutions to demonstrate how TF1's `tf.estimator`-derived models can be migrated to TF2 with Keras.\n",
        "\n",
        "Namely, this guide includes examples for migrating:\n",
        "*   `tf.estimator.LinearEstimator`/`Classifier`/`Regressor` -> LinearModel (Keras model defined in this guide)\n",
        "*   `tf.estimator.DNNEstimator`/`Classifier`/`Regressor` -> Custom Keras DNN Model\n",
        "*   `tf.estimator.DNNLinearCombinedEstimator`/`Classifier`/`Regressor` -> WideDeepModel  (Keras model defined in this guide)\n",
        "*   `tf.estimator.BoostedTreesEstimator`/`Classifier`/`Regressor` -> `tfdf.keras.GradientBoostedTreesModel`, `tfdf.keras.RandomForestModel`\n",
        "\n",
        "A common precursor to the training of a model is feature preprocessing, which is done for TF1 Estimator models with `tf.feature_column`. For more information on feature preprocessing in TF2, see [this guide on migrating feature columns](migrating_feature_columns.ipynb).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdZSoIXEbhg-"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Start with a couple of necessary TensorFlow imports,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qsgZp0f-nu9s"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_decision_forests"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iE0vSfMXumKI"
      },
      "outputs": [],
      "source": [
        "import keras\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1\n",
        "import tensorflow_decision_forests as tfdf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jsm9Rxx7s1OZ"
      },
      "source": [
        "prepare some simple data for demonstration from the standard Titanic dataset,"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wC6i_bEZPrPY"
      },
      "outputs": [],
      "source": [
        "x_train = pd.read_csv('https://storage.googleapis.com/tf-datasets/titanic/train.csv')\n",
        "x_eval = pd.read_csv('https://storage.googleapis.com/tf-datasets/titanic/eval.csv')\n",
        "x_train['sex'].replace(('male', 'female'), (0, 1), inplace=True)\n",
        "x_eval['sex'].replace(('male', 'female'), (0, 1), inplace=True)\n",
        "\n",
        "x_train['alone'].replace(('n', 'y'), (0, 1), inplace=True)\n",
        "x_eval['alone'].replace(('n', 'y'), (0, 1), inplace=True)\n",
        "\n",
        "x_train['class'].replace(('First', 'Second', 'Third'), (1, 2, 3), inplace=True)\n",
        "x_eval['class'].replace(('First', 'Second', 'Third'), (1, 2, 3), inplace=True)\n",
        "\n",
        "x_train.drop(['embark_town', 'deck'], axis=1, inplace=True)\n",
        "x_eval.drop(['embark_town', 'deck'], axis=1, inplace=True)\n",
        "\n",
        "y_train = x_train.pop('survived')\n",
        "y_eval = x_eval.pop('survived')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lqe9obf7suIj"
      },
      "outputs": [],
      "source": [
        "## DATA SETUP FOR ESTIMATORS (TENSORFLOW 1)\n",
        "def _input_fn():\n",
        "  return tf1.data.Dataset.from_tensor_slices((dict(x_train), y_train)).batch(32)\n",
        "\n",
        "\n",
        "def _eval_input_fn():\n",
        "  return tf1.data.Dataset.from_tensor_slices((dict(x_eval), y_eval)).batch(32)\n",
        "\n",
        "\n",
        "FEATURE_NAMES = [\n",
        "    'age', 'fare', 'sex', 'n_siblings_spouses', 'parch', 'class', 'alone'\n",
        "]\n",
        "\n",
        "feature_columns = []\n",
        "for fn in FEATURE_NAMES:\n",
        "  feat_col = tf1.feature_column.numeric_column(fn, dtype=tf.float32)\n",
        "  feature_columns.append(feat_col)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bYSgoezeMrpI"
      },
      "source": [
        "and create a method to instantiate a simplistic sample optimizer to use with our various TF1 Estimator and TF2 Keras models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YHB_nuzVLVLe"
      },
      "outputs": [],
      "source": [
        "def create_sample_optimizer(tf_version):\n",
        "  if tf_version == 'tf1':\n",
        "    optimizer = lambda: tf.keras.optimizers.Ftrl(\n",
        "        l1_regularization_strength=0.001,\n",
        "        learning_rate=tf1.train.exponential_decay(\n",
        "            learning_rate=0.1,\n",
        "            global_step=tf1.train.get_global_step(),\n",
        "            decay_steps=10000,\n",
        "            decay_rate=0.9))\n",
        "  elif tf_version == 'tf2':\n",
        "    optimizer = tf.keras.optimizers.Ftrl(\n",
        "        l1_regularization_strength=0.001,\n",
        "        learning_rate=tf.keras.optimizers.schedules.ExponentialDecay(\n",
        "            initial_learning_rate=0.1, decay_steps=10000, decay_rate=0.9))\n",
        "  return optimizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4uXff1BEssdE"
      },
      "source": [
        "## Example 1: Migrating from LinearEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_O7fyhCnpvED"
      },
      "source": [
        "### TF1: Using LinearEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A9560BqEOTpb"
      },
      "source": [
        "In TF1, you can use `tf.estimator.LinearEstimator` to create a baseline linear model for regression and classification problems."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oWfh0QW4IXTn"
      },
      "outputs": [],
      "source": [
        "linear_estimator = tf.estimator.LinearEstimator(\n",
        "    head=tf.estimator.BinaryClassHead(),\n",
        "    feature_columns=feature_columns,\n",
        "    optimizer=create_sample_optimizer('tf1'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hi77Sg4k-0TR"
      },
      "outputs": [],
      "source": [
        "linear_estimator.train(input_fn=_input_fn, steps=100)\n",
        "linear_estimator.evaluate(input_fn=_eval_input_fn, steps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEmzBjfnsxwT"
      },
      "source": [
        "### TF2: Using Keras LinearModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fkgkGf_AOaRR"
      },
      "source": [
        "In TF2, you can create an instance of the Keras `LinearModel` implementation below to substitute for one generated by `tf.estimator.LinearEstimator`, with similar levels of user-specified customization (for example, the ability to customize a chosen model optimizer)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iwHfyVSxL5No"
      },
      "outputs": [],
      "source": [
        "class LinearModel(keras.engine.training.Model):\n",
        "  r\"\"\"Linear Model for regression and classification problems.\n",
        "\n",
        "  This model approximates the following function:\n",
        "  $$y = \\beta + \\sum_{i=1}^{N} w_{i} * x_{i}$$\n",
        "  where $$\\beta$$ is the bias and $$w_{i}$$ is the weight for each feature.\n",
        "\n",
        "  Example:\n",
        "\n",
        "  ```python\n",
        "  model = LinearModel()\n",
        "  model.compile(optimizer='sgd', loss='mse')\n",
        "  model.fit(x, y, epochs=epochs)\n",
        "  ```\n",
        "\n",
        "  This model accepts sparse float inputs as well:\n",
        "\n",
        "  Example:\n",
        "  ```python\n",
        "  model = LinearModel()\n",
        "  opt = tf.keras.optimizers.Adam()\n",
        "  loss_fn = tf.keras.losses.MeanSquaredError()\n",
        "  with tf.GradientTape() as tape:\n",
        "    output = model(sparse_input)\n",
        "    loss = tf.reduce_mean(loss_fn(target, output))\n",
        "  grads = tape.gradient(loss, model.weights)\n",
        "  opt.apply_gradients(zip(grads, model.weights))\n",
        "  ```\n",
        "\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               units=1,\n",
        "               activation=None,\n",
        "               use_bias=True,\n",
        "               kernel_initializer='zeros',\n",
        "               bias_initializer='zeros',\n",
        "               kernel_regularizer=None,\n",
        "               bias_regularizer=None,\n",
        "               **kwargs):\n",
        "    \"\"\"Create a Linear Model.\n",
        "\n",
        "    Args:\n",
        "      units: Positive integer, output dimension without the batch size.\n",
        "      activation: Activation function to use.\n",
        "        If you don't specify anything, no activation is applied.\n",
        "      use_bias: whether to calculate the bias/intercept for this model. If set\n",
        "        to False, no bias/intercept will be used in calculations, e.g., the data\n",
        "        is already centered.\n",
        "      kernel_initializer: Initializer for the `kernel` weights matrices.\n",
        "      bias_initializer: Initializer for the bias vector.\n",
        "      kernel_regularizer: regularizer for kernel vectors.\n",
        "      bias_regularizer: regularizer for bias vector.\n",
        "      **kwargs: The keyword arguments that are passed on to BaseLayer.__init__.\n",
        "    \"\"\"\n",
        "\n",
        "    self.units = units\n",
        "    self.activation = keras.activations.get(activation)\n",
        "    self.use_bias = use_bias\n",
        "    self.kernel_initializer = keras.initializers.get(kernel_initializer)\n",
        "    self.bias_initializer = keras.initializers.get(bias_initializer)\n",
        "    self.kernel_regularizer = keras.regularizers.get(kernel_regularizer)\n",
        "    self.bias_regularizer = keras.regularizers.get(bias_regularizer)\n",
        "    super(LinearModel, self).__init__(**kwargs)\n",
        "    keras.engine.base_layer.keras_premade_model_gauge.get_cell('Linear').set(True)\n",
        "\n",
        "  def build(self, input_shape):\n",
        "    if isinstance(input_shape, dict):\n",
        "      names = sorted(list(input_shape.keys()))\n",
        "      self.input_specs = []\n",
        "      self.dense_layers = []\n",
        "      for name in names:\n",
        "        shape = input_shape[name]\n",
        "        layer = keras.layers.core.Dense(\n",
        "            units=self.units,\n",
        "            use_bias=False,\n",
        "            kernel_regularizer=self.kernel_regularizer,\n",
        "            name=name)\n",
        "        self.input_specs.append(\n",
        "            keras.engine.input_spec.InputSpec(shape=shape, name=name))\n",
        "        self.dense_layers.append(layer)\n",
        "    elif isinstance(input_shape, (tuple, list)) and all(\n",
        "        isinstance(shape, tf.TensorShape) for shape in input_shape):\n",
        "      self.dense_layers = []\n",
        "      for shape in input_shape:\n",
        "        layer = keras.layers.core.Dense(\n",
        "            units=self.units,\n",
        "            use_bias=False,\n",
        "            kernel_initializer=self.kernel_initializer,\n",
        "            kernel_regularizer=self.kernel_regularizer)\n",
        "        layer.build(shape)\n",
        "        self.dense_layers.append(layer)\n",
        "    else:\n",
        "      # input_shape can be a single TensorShape or a tuple of ints.\n",
        "      layer = keras.layers.core.Dense(\n",
        "          units=self.units,\n",
        "          use_bias=False,\n",
        "          kernel_initializer=self.kernel_initializer,\n",
        "          kernel_regularizer=self.kernel_regularizer)\n",
        "      layer.build(input_shape)\n",
        "      self.dense_layers = [layer]\n",
        "\n",
        "    if self.use_bias:\n",
        "      self.bias = self.add_weight(\n",
        "          'bias',\n",
        "          shape=self.units,\n",
        "          initializer=self.bias_initializer,\n",
        "          regularizer=self.bias_regularizer,\n",
        "          dtype=self.dtype,\n",
        "          trainable=True)\n",
        "    else:\n",
        "      self.bias = None\n",
        "    self.built = True\n",
        "\n",
        "  def call(self, inputs):\n",
        "    result = None\n",
        "    if isinstance(inputs, dict):\n",
        "      names = [layer.name for layer in self.dense_layers]\n",
        "      different_keys = set(names) - set(inputs.keys())\n",
        "      if different_keys:\n",
        "        raise ValueError(\n",
        "            'The input dictionary does not match '\n",
        "            'the structure expected by the model.'\n",
        "            '\\n\\tExpected keys: {}'\n",
        "            '\\n\\tReceived keys: {}'\n",
        "            '\\n\\tMissing keys: {}'.format(set(names), set(inputs.keys()),\n",
        "                                          different_keys))\n",
        "      inputs = [inputs[name] for name in names]\n",
        "      for inp, layer in zip(inputs, self.dense_layers):\n",
        "        output = layer(inp)\n",
        "        if result is None:\n",
        "          result = output\n",
        "        else:\n",
        "          result += output\n",
        "    elif isinstance(inputs, (tuple, list)):\n",
        "      for inp, layer in zip(inputs, self.dense_layers):\n",
        "        output = layer(inp)\n",
        "        if result is None:\n",
        "          result = output\n",
        "        else:\n",
        "          result += output\n",
        "    else:\n",
        "      result = self.dense_layers[0](inputs)\n",
        "\n",
        "    if self.use_bias:\n",
        "      result = tf.nn.bias_add(result, self.bias)\n",
        "    if self.activation is not None:\n",
        "      return self.activation(result)  # pylint: disable=not-callable\n",
        "    return result\n",
        "\n",
        "  def get_config(self):\n",
        "    config = {\n",
        "        'units': self.units,\n",
        "        'activation': keras.activations.serialize(self.activation),\n",
        "        'use_bias': self.use_bias,\n",
        "        'kernel_initializer': keras.initializers.serialize(self.kernel_initializer),\n",
        "        'bias_initializer': keras.initializers.serialize(self.bias_initializer),\n",
        "        'kernel_regularizer': keras.regularizers.serialize(self.kernel_regularizer),\n",
        "        'bias_regularizer': keras.regularizers.serialize(self.bias_regularizer),\n",
        "    }\n",
        "    base_config = keras.engine.base_layer.Layer.get_config(self)\n",
        "    return dict(list(base_config.items()) + list(config.items()))\n",
        "\n",
        "  @classmethod\n",
        "  def from_config(cls, config, custom_objects=None):\n",
        "    del custom_objects\n",
        "    return cls(**config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kip65sYBlKiu"
      },
      "outputs": [],
      "source": [
        "linear_model = LinearModel()\n",
        "linear_model.compile(loss='mse', optimizer=create_sample_optimizer('tf2'), metrics=['accuracy'])\n",
        "linear_model.fit(x_train, y_train, epochs=10)\n",
        "linear_model.evaluate(x_eval, y_eval, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RRrj78Lqplni"
      },
      "source": [
        "## Example 2: Migrating from DNNEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YKl6XZ7Bp1t5"
      },
      "source": [
        "### TF1: Using DNNEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J7wJUmgypln8"
      },
      "source": [
        "In TF1, you can use `tf.estimator.DNNEstimator` to create a baseline DNN model for regression and classification problems."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qHbgXCzfpln9"
      },
      "outputs": [],
      "source": [
        "dnn_estimator = tf.estimator.DNNEstimator(\n",
        "    head=tf.estimator.BinaryClassHead(),\n",
        "    feature_columns=feature_columns,\n",
        "    hidden_units=[128],\n",
        "    activation_fn=tf.nn.relu,\n",
        "    optimizer=create_sample_optimizer('tf1'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6DTnXxU2pln-"
      },
      "outputs": [],
      "source": [
        "dnn_estimator.train(input_fn=_input_fn, steps=100)\n",
        "dnn_estimator.evaluate(input_fn=_eval_input_fn, steps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xJz6px6pln-"
      },
      "source": [
        "### TF2: Using Keras to Create a Custom DNN Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7cgc72rzpln-"
      },
      "source": [
        "In TF2, you can create a custom DNN model to substitute for one generated by `tf.estimator.DNNEstimator`, with similar levels of user-specified customization (for instance, as in the previous example, the ability to customize a chosen model optimizer).\n",
        "\n",
        "A similar workflow can be used to replace `tf.estimator.experimental.RNNEstimator` with a Keras RNN Model. Keras provides a number of built-in, customizable choices by way of `tf.keras.layers.RNN`, `tf.keras.layers.LSTM`, and `tf.keras.layers.GRU` - see [here](https://www.tensorflow.org/guide/keras/rnn#built-in_rnn_layers_a_simple_example) for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B5SdsjlL49RG"
      },
      "outputs": [],
      "source": [
        "dnn_model = tf.keras.models.Sequential(\n",
        "    [tf.keras.layers.Dense(128, activation='relu'),\n",
        "     tf.keras.layers.Dense(1)])\n",
        "\n",
        "dnn_model.compile(loss='mse', optimizer=create_sample_optimizer('tf2'), metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JQmRw9_Upln_"
      },
      "outputs": [],
      "source": [
        "dnn_model.fit(x_train, y_train, epochs=10)\n",
        "dnn_model.evaluate(x_eval, y_eval, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UeBHZ0cd1Pl2"
      },
      "source": [
        "## Example 3: Migrating from DNNLinearCombinedEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GfRaObf5g4TU"
      },
      "source": [
        "### TF1: Using DNNLinearCombinedEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2r13RMX-g4TV"
      },
      "source": [
        "In TF1, you can use `tf.estimator.DNNLinearCombinedEstimator` to create a baseline combined model for regression and classification problems with customization capacity for both its linear and DNN components."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OyyDCqc5j7rf"
      },
      "outputs": [],
      "source": [
        "optimizer = create_sample_optimizer('tf1')\n",
        "\n",
        "combined_estimator = tf.estimator.DNNLinearCombinedEstimator(\n",
        "    head=tf.estimator.BinaryClassHead(),\n",
        "    # Wide settings\n",
        "    linear_feature_columns=feature_columns,\n",
        "    linear_optimizer=optimizer,\n",
        "    # Deep settings\n",
        "    dnn_feature_columns=feature_columns,\n",
        "    dnn_hidden_units=[128],\n",
        "    dnn_optimizer=optimizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aXN-BxwzmRaf"
      },
      "outputs": [],
      "source": [
        "combined_estimator.train(input_fn=_input_fn, steps=100)\n",
        "combined_estimator.evaluate(input_fn=_eval_input_fn, steps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BeMikL5ug4TX"
      },
      "source": [
        "### TF2: Using Keras WideDeepModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CYByxxBhg4TX"
      },
      "source": [
        "In TF2, you can create an instance of the Keras `WideDeepModel` implementation below to substitute for one generated by `tf.estimator.DNNLinearCombinedEstimator`, with similar levels of user-specified customization (for instance, as in the previous example, the ability to customize a chosen model optimizer).\n",
        "\n",
        "This `WideDeepModel` is constructed on the basis of a constituent `LinearModel` and a custom DNN Model, both of which are created in the preceding two examples. A custom linear model can also be used in place of the built-in Keras `LinearModel` if desired."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sYubfpmuP3Hi"
      },
      "outputs": [],
      "source": [
        "class WideDeepModel(keras.engine.training.Model):\n",
        "  r\"\"\"Wide & Deep Model for regression and classification problems.\n",
        "\n",
        "  This model jointly train a linear and a dnn model.\n",
        "\n",
        "  Example:\n",
        "\n",
        "  ```python\n",
        "  linear_model = LinearModel()\n",
        "  dnn_model = keras.Sequential([keras.layers.Dense(units=64),\n",
        "                               keras.layers.Dense(units=1)])\n",
        "  combined_model = WideDeepModel(linear_model, dnn_model)\n",
        "  combined_model.compile(optimizer=['sgd', 'adam'], 'mse', ['mse'])\n",
        "  # define dnn_inputs and linear_inputs as separate numpy arrays or\n",
        "  # a single numpy array if dnn_inputs is same as linear_inputs.\n",
        "  combined_model.fit([linear_inputs, dnn_inputs], y, epochs)\n",
        "  # or define a single `tf.data.Dataset` that contains a single tensor or\n",
        "  # separate tensors for dnn_inputs and linear_inputs.\n",
        "  dataset = tf.data.Dataset.from_tensors(([linear_inputs, dnn_inputs], y))\n",
        "  combined_model.fit(dataset, epochs)\n",
        "  ```\n",
        "\n",
        "  Both linear and dnn model can be pre-compiled and trained separately\n",
        "  before jointly training:\n",
        "\n",
        "  Example:\n",
        "  ```python\n",
        "  linear_model = LinearModel()\n",
        "  linear_model.compile('adagrad', 'mse')\n",
        "  linear_model.fit(linear_inputs, y, epochs)\n",
        "  dnn_model = keras.Sequential([keras.layers.Dense(units=1)])\n",
        "  dnn_model.compile('rmsprop', 'mse')\n",
        "  dnn_model.fit(dnn_inputs, y, epochs)\n",
        "  combined_model = WideDeepModel(linear_model, dnn_model)\n",
        "  combined_model.compile(optimizer=['sgd', 'adam'], 'mse', ['mse'])\n",
        "  combined_model.fit([linear_inputs, dnn_inputs], y, epochs)\n",
        "  ```\n",
        "\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, linear_model, dnn_model, activation=None, **kwargs):\n",
        "    \"\"\"Create a Wide & Deep Model.\n",
        "\n",
        "    Args:\n",
        "      linear_model: a premade LinearModel, its output must match the output of\n",
        "        the dnn model.\n",
        "      dnn_model: a `tf.keras.Model`, its output must match the output of the\n",
        "        linear model.\n",
        "      activation: Activation function. Set it to None to maintain a linear\n",
        "        activation.\n",
        "      **kwargs: The keyword arguments that are passed on to BaseLayer.__init__.\n",
        "        Allowed keyword arguments include `name`.\n",
        "    \"\"\"\n",
        "    super(WideDeepModel, self).__init__(**kwargs)\n",
        "    keras.engine.base_layer.keras_premade_model_gauge.get_cell('WideDeep').set(True)\n",
        "    self.linear_model = linear_model\n",
        "    self.dnn_model = dnn_model\n",
        "    self.activation = keras.activations.get(activation)\n",
        "\n",
        "  def call(self, inputs, training=None):\n",
        "    if not isinstance(inputs, (tuple, list)) or len(inputs) != 2:\n",
        "      linear_inputs = dnn_inputs = inputs\n",
        "    else:\n",
        "      linear_inputs, dnn_inputs = inputs\n",
        "    linear_output = self.linear_model(linear_inputs)\n",
        "    # pylint: disable=protected-access\n",
        "    if self.dnn_model._expects_training_arg:\n",
        "      if training is None:\n",
        "        training = keras.backend.learning_phase()\n",
        "      dnn_output = self.dnn_model(dnn_inputs, training=training)\n",
        "    else:\n",
        "      dnn_output = self.dnn_model(dnn_inputs)\n",
        "    output = tf.nest.map_structure(lambda x, y: (x + y), linear_output,\n",
        "                                   dnn_output)\n",
        "    if self.activation:\n",
        "      return tf.nest.map_structure(self.activation, output)\n",
        "    return output\n",
        "\n",
        "  # This does not support gradient scaling and LossScaleOptimizer.\n",
        "  def train_step(self, data):\n",
        "    x, y, sample_weight = keras.engine.data_adapter.unpack_x_y_sample_weight(data)\n",
        "\n",
        "    with tf.GradientTape() as tape:\n",
        "      y_pred = self(x, training=True)\n",
        "      loss = self.compiled_loss(\n",
        "          y, y_pred, sample_weight, regularization_losses=self.losses)\n",
        "    self.compiled_metrics.update_state(y, y_pred, sample_weight)\n",
        "\n",
        "    if isinstance(self.optimizer, (list, tuple)):\n",
        "      linear_vars = self.linear_model.trainable_variables\n",
        "      dnn_vars = self.dnn_model.trainable_variables\n",
        "      linear_grads, dnn_grads = tape.gradient(loss, (linear_vars, dnn_vars))\n",
        "\n",
        "      linear_optimizer = self.optimizer[0]\n",
        "      dnn_optimizer = self.optimizer[1]\n",
        "      linear_optimizer.apply_gradients(zip(linear_grads, linear_vars))\n",
        "      dnn_optimizer.apply_gradients(zip(dnn_grads, dnn_vars))\n",
        "    else:\n",
        "      trainable_variables = self.trainable_variables\n",
        "      grads = tape.gradient(loss, trainable_variables)\n",
        "      self.optimizer.apply_gradients(zip(grads, trainable_variables))\n",
        "\n",
        "    return {m.name: m.result() for m in self.metrics}\n",
        "\n",
        "  def _make_train_function(self):\n",
        "    # Only needed for graph mode and model_to_estimator.\n",
        "    has_recompiled = self._recompile_weights_loss_and_weighted_metrics()\n",
        "    self._check_trainable_weights_consistency()\n",
        "    # If the loss/weighted metric sub-graphs have been re-compiled then create\n",
        "    # train function even if one exists already. This is because\n",
        "    # `_feed_sample_weights` list has been updated on re-compile.\n",
        "    if getattr(self, 'train_function', None) is None or has_recompiled:\n",
        "      # Restore the compiled trainable state.\n",
        "      current_trainable_state = self._get_trainable_state()\n",
        "      self._set_trainable_state(self._compiled_trainable_state)\n",
        "\n",
        "      inputs = (\n",
        "          self._feed_inputs + self._feed_targets + self._feed_sample_weights)\n",
        "      if not isinstance(keras.backend.symbolic_learning_phase(), int):\n",
        "        inputs += [keras.backend.symbolic_learning_phase()]\n",
        "\n",
        "      if isinstance(self.optimizer, (list, tuple)):\n",
        "        linear_optimizer = self.optimizer[0]\n",
        "        dnn_optimizer = self.optimizer[1]\n",
        "      else:\n",
        "        linear_optimizer = self.optimizer\n",
        "        dnn_optimizer = self.optimizer\n",
        "\n",
        "      with keras.backend.get_graph().as_default():\n",
        "        with keras.backend.name_scope('training'):\n",
        "          # Training updates\n",
        "          updates = []\n",
        "          linear_updates = linear_optimizer.get_updates(\n",
        "              params=self.linear_model.trainable_weights,  # pylint: disable=protected-access\n",
        "              loss=self.total_loss)\n",
        "          updates += linear_updates\n",
        "          dnn_updates = dnn_optimizer.get_updates(\n",
        "              params=self.dnn_model.trainable_weights,  # pylint: disable=protected-access\n",
        "              loss=self.total_loss)\n",
        "          updates += dnn_updates\n",
        "          # Unconditional updates\n",
        "          updates += self.get_updates_for(None)\n",
        "          # Conditional updates relevant to this model\n",
        "          updates += self.get_updates_for(self.inputs)\n",
        "\n",
        "        metrics = self._get_training_eval_metrics()\n",
        "        metrics_tensors = [\n",
        "            m._call_result for m in metrics if hasattr(m, '_call_result')  # pylint: disable=protected-access\n",
        "        ]\n",
        "\n",
        "      with keras.backend.name_scope('training'):\n",
        "        # Gets loss and metrics. Updates weights at each call.\n",
        "        fn = keras.backend.function(\n",
        "            inputs, [self.total_loss] + metrics_tensors,\n",
        "            updates=updates,\n",
        "            name='train_function',\n",
        "            **self._function_kwargs)\n",
        "        setattr(self, 'train_function', fn)\n",
        "\n",
        "      # Restore the current trainable state\n",
        "      self._set_trainable_state(current_trainable_state)\n",
        "\n",
        "  def get_config(self):\n",
        "    linear_config = keras.utils.serialize_keras_object(self.linear_model)\n",
        "    dnn_config = keras.utils.serialize_keras_object(self.dnn_model)\n",
        "    config = {\n",
        "        'linear_model': linear_config,\n",
        "        'dnn_model': dnn_config,\n",
        "        'activation': keras.activations.serialize(self.activation),\n",
        "    }\n",
        "    base_config = keras.engine.base_layer.Layer.get_config(self)\n",
        "    return dict(list(base_config.items()) + list(config.items()))\n",
        "\n",
        "  @classmethod\n",
        "  def from_config(cls, config, custom_objects=None):\n",
        "    linear_config = config.pop('linear_model')\n",
        "    linear_model = keras.layers.deserialize(linear_config, custom_objects)\n",
        "    dnn_config = config.pop('dnn_model')\n",
        "    dnn_model = keras.layers.deserialize(dnn_config, custom_objects)\n",
        "    activation = keras.activations.deserialize(\n",
        "        config.pop('activation', None), custom_objects=custom_objects)\n",
        "    return cls(\n",
        "        linear_model=linear_model,\n",
        "        dnn_model=dnn_model,\n",
        "        activation=activation,\n",
        "        **config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mIFM3e-_RLSX"
      },
      "outputs": [],
      "source": [
        "# Create LinearModel and DNN Model as in Examples 1 and 2\n",
        "optimizer = create_sample_optimizer('tf2')\n",
        "\n",
        "linear_model = LinearModel()\n",
        "linear_model.compile(loss='mse', optimizer=optimizer, metrics=['accuracy'])\n",
        "linear_model.fit(x_train, y_train, epochs=10, verbose=0)\n",
        "\n",
        "dnn_model = tf.keras.models.Sequential(\n",
        "    [tf.keras.layers.Dense(128, activation='relu'),\n",
        "     tf.keras.layers.Dense(1)])\n",
        "dnn_model.compile(loss='mse', optimizer=optimizer, metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mFmQz9kjmMSx"
      },
      "outputs": [],
      "source": [
        "combined_model = WideDeepModel(linear_model, dnn_model)\n",
        "combined_model.compile(\n",
        "    optimizer=[optimizer, optimizer], loss='mse', metrics=['accuracy'])\n",
        "combined_model.fit([x_train, x_train], y_train, epochs=10)\n",
        "combined_model.evaluate(x_eval, y_eval, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wP1DBRhpeOJn"
      },
      "source": [
        "## Example 4: Migrating from BoostedTreesEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_3mCQVDSeOKD"
      },
      "source": [
        "### TF1: Using BoostedTreesEstimator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oEWYHNt4eOKD"
      },
      "source": [
        "In TF1, you can use `tf.estimator.BoostedTreesEstimator` to create a baseline to create a baseline Gradient Boosting model using an ensemble of decision trees for regression and classification problems."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uBIURLLMeOKE"
      },
      "outputs": [],
      "source": [
        "bt_estimator = tf1.estimator.BoostedTreesEstimator(\n",
        "    head=tf.estimator.BinaryClassHead(),\n",
        "    n_batches_per_layer=1,\n",
        "    max_depth=10,\n",
        "    n_trees=1000,\n",
        "    feature_columns=feature_columns)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MUo5omt3eOKE"
      },
      "outputs": [],
      "source": [
        "bt_estimator.train(input_fn=_input_fn, steps=1000)\n",
        "bt_estimator.evaluate(input_fn=_eval_input_fn, steps=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eNuLP6BeeOKF"
      },
      "source": [
        "### TF2: Using TensorFlow Decision Forests"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m3EVq388eOKF"
      },
      "source": [
        "In TF2, the closest pre-packaged substitute for a model generated by `tf.estimator.BoostedTreesEstimator` is one created using `tfdf.keras.GradientBoostedTreesModel`, which creates a sequentially-trained sequence of shallow decision trees, each designed to \"learn\" from errors made by its predecessors in the sequence.\n",
        "\n",
        "`GradientBoostedTreesModel` provides more options for customization, allowing for the specification of everything from basic depth constraints to early stopping conditions. See [here](https://www.tensorflow.org/decision_forests/api_docs/python/tfdf/keras/GradientBoostedTreesModel#attributes) for more `GradientBoostedTreesModel` attribute details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JLS_2vKKeOKF"
      },
      "outputs": [],
      "source": [
        "gbt_model = tfdf.keras.GradientBoostedTreesModel(\n",
        "    task=tfdf.keras.Task.CLASSIFICATION)\n",
        "gbt_model.compile(metrics=['mse', 'accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sZZSM7_VeOKF"
      },
      "outputs": [],
      "source": [
        "train_df, eval_df = x_train.copy(), x_eval.copy()\n",
        "train_df['survived'], eval_df['survived'] = y_train, y_eval\n",
        "\n",
        "train_dataset = tfdf.keras.pd_dataframe_to_tf_dataset(train_df, label='survived')\n",
        "eval_dataset = tfdf.keras.pd_dataframe_to_tf_dataset(eval_df, label='survived')\n",
        "\n",
        "gbt_model.fit(train_dataset)\n",
        "gbt_model.evaluate(eval_dataset, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z22UJ5SUqToQ"
      },
      "source": [
        "In TF2, there is also another available TFDF substitute for a model generated by `tf.estimator.BoostedTreesEstimator` - `tfdf.keras.RandomForestModel`. `RandomForestModel` creates a robust, overfitting-resistant learner consisting of a voting population of deep decision trees, each trained on random subsets of the input training dataset.\n",
        "\n",
        "`RandomForestModel` and `GradientBoostedTreesModel` provide similarly extensive levels of customization. Choosing between them is problem-specific and dependent on your task or application.\n",
        "\n",
        "See [here](https://www.tensorflow.org/decision_forests/api_docs/python/tfdf/keras/GradientBoostedTreesModel#attributes) for more `RandomForestModel` attribute details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "027bGnCork_W"
      },
      "outputs": [],
      "source": [
        "rf_model = tfdf.keras.RandomForestModel(\n",
        "    task=tfdf.keras.Task.CLASSIFICATION)\n",
        "rf_model.compile(metrics=['mse', 'accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tc7KtjMlryn_"
      },
      "outputs": [],
      "source": [
        "rf_model.fit(train_dataset)\n",
        "rf_model.evaluate(eval_dataset, return_dict=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "canned_estimators.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
